In [1]:
import os
import time
import numpy as np
import scipy as sp
import scipy.io as sio
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy.linalg as la
import glob
import nodepy.linear_multistep_method as lm
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random
import tensorflow as tf
import tensorflow_datasets as tfds
#import haiku as hk
import jax
import flax.linen as nn
from flax.training import train_state # Useful dataclass to keep train state
from typing import Sequence
import optax
from jax import jacfwd, jacrev
In [2]:
import jax
jax.__version__
Out[2]:
In [3]:
import seaborn as sns
from matplotlib import rc
def set_style_sns():
sns.set_context('paper')
sns.set(font = 'serif')
sns.set(font_scale = 1.3)
sns.set_style('white', {
'font.family': 'serif',
'font.serif': ['Time', 'Palatino', 'serif'],
'lines.markersize': 10
})
plt.rcParams.update({'font.size':16})
set_style_sns()
plt.rcParams.update({'text.usetex': False})
In [4]:
import scipy.io as sio
data = sio.loadmat('DOF6_ROM_ml-20230324-124511cubic_conserve-1-dof0064')
In [5]:
Mhat = data['Mhat']
Khat = data['Khat']
Chat = np.zeros((6,6)) #data['Chat']
t = data['t']
dt = data['dt']
q_ = data['Qr']
dq_ = data['Qdr']
ddq_ = data['Qddr']
x = data['Q']
dx = data['Qd']
ddx = data['Qdd']
V = data['v']
q = np.dot(V.T, x)
dq = np.dot(V.T, dx)
ddq = np.dot(V.T, ddx)
Qext = np.zeros(q.shape)
Qrr = data['Qrr']
x_ = np.dot(V, q)
dx_ = np.dot(V, dq)
ddx_ = np.dot(V, ddq)
Ntt = data['Ntt'][0][0]
print(q.shape, dq.shape, ddq.shape, q_.shape)
In [6]:
Ntt
Out[6]:
In [7]:
print(Khat.shape, Ntt)
In [8]:
plt.figure(figsize = (10,4))
sample = 10
plt.plot(t[0], x.T[:, sample], 'b', alpha = 0.5)
plt.plot(t[0], Qrr.T[2:, sample], 'r', alpha = 0.2)
plt.ylabel('Displacement')
plt.xlabel('Time, s')
plt.show()
In [9]:
plt.figure(figsize = (10,4))
sample = np.array([33, 10])
plt.plot(t[0], ddx.T[:, sample])
plt.plot(t[0],ddx_.T[:, sample], 'k--')
plt.xlim([0,18])
plt.show()
In [10]:
plt.figure(figsize = (14,6))
#Ntt = q.shape[1]
plt.subplot(4,1,1)
plt.plot(t[0,:Ntt], q[:,:Ntt].T)
plt.subplot(4,1,2)
plt.plot(t[0,:Ntt], dq[:,:Ntt].T)
plt.subplot(4,1,3)
plt.plot(t[0,:Ntt], ddq[:,:Ntt].T)
plt.tight_layout()
plt.show()
In [11]:
Ntt
Out[11]:
In [54]:
nDof = q.shape[0]
test_factor = 2
start = 0
subset = Ntt*test_factor #8200 #45000 #12000
ixend = Ntt*test_factor + 200 #8600 #6500 #45500 #12500
# Choose training data - using linear
QnormTrainData = np.concatenate([q.T[start:,:,np.newaxis], dq.T[start:,:,np.newaxis], ddq.T[start:,:,np.newaxis], Qext.T[start:,:,np.newaxis]], axis = 2).reshape((-1,nDof,4))
#xmax_for_norm = np.max(np.max(normTrainData, axis = 0), axis = 0)
#normTrainData = normTrainData/xmax
QnormTrainData = QnormTrainData[:subset, :, :]
q_train_dataset = tf.data.Dataset.from_tensor_slices(tf.cast(QnormTrainData,dtype=tf.float32))
#train_dataset = train_dataset.batch(256)
QnormTestData = np.concatenate([q.T[start:,:,np.newaxis], dq.T[start:,:,np.newaxis], ddq.T[start:,:,np.newaxis], Qext.T[start:,:,np.newaxis]], axis = 2).reshape((-1,nDof,4))
QnormTestData = QnormTestData[subset:ixend, :, :]
q_test_dataset = tf.data.Dataset.from_tensor_slices(tf.cast(QnormTestData,dtype=tf.float32))
#test_dataset = test_dataset.batch(512)
In [55]:
# Choose training data - using linear
Fext = np.zeros(x.shape)
full_nDof = x.shape[0]
XnormTrainData = np.concatenate([x.T[start:,:,np.newaxis], dx.T[start:,:,np.newaxis], ddx.T[start:,:,np.newaxis], Fext.T[start:,:,np.newaxis]], axis = 2).reshape((-1,full_nDof,4))
#xmax_for_norm = np.max(np.max(normTrainData, axis = 0), axis = 0)
#normTrainData = normTrainData/xmax
XnormTrainData = XnormTrainData[:subset, :, :]
x_train_dataset = tf.data.Dataset.from_tensor_slices(tf.cast(XnormTrainData,dtype=tf.float32))
#train_dataset = train_dataset.batch(256)
XnormTestData = np.concatenate([x.T[start:,:,np.newaxis], dx.T[start:,:,np.newaxis], ddx.T[start:,:,np.newaxis], Fext.T[start:,:,np.newaxis]], axis = 2).reshape((-1,full_nDof,4))
XnormTestData = XnormTestData[subset:ixend, :, :]
x_test_dataset = tf.data.Dataset.from_tensor_slices(tf.cast(XnormTestData,dtype=tf.float32))
#test_dataset = test_dataset.batch(512)
In [15]:
print(XnormTrainData.shape, XnormTestData.shape)
print(QnormTrainData.shape, QnormTestData.shape)
In [16]:
plt.figure(figsize = (10,6))
plt.subplot(2,1,1)
plt.plot(QnormTrainData[:,:,1])
plt.plot(np.diff(QnormTrainData[:,:,0], axis = 0)/dt, 'k--')
plt.subplot(2,1,2)
plt.plot(QnormTrainData[:,:,2])
plt.plot(np.diff(QnormTrainData[:,:,1], axis = 0)/dt, 'k--')
plt.show()
In [17]:
plt.figure(figsize = (10,6))
plt.subplot(2,1,1)
plt.plot(XnormTrainData[:,:,1])
plt.plot(np.diff(XnormTrainData[:,:,0], axis = 0)/dt, 'k--')
plt.xlim([1000,1200])
plt.subplot(2,1,2)
plt.plot(XnormTrainData[:,:,2])
plt.plot(np.diff(XnormTrainData[:,:,1], axis = 0)/dt, 'k--')
plt.xlim([1000,1200])
plt.show()
In [18]:
NLAYERS = 4
NNODES = 6 #10
class MLP(nn.Module):
features: Sequence[int]
@nn.compact
def __call__(self, xx):
for feat in self.features[:-1]:
xx = nn.tanh(nn.Dense(feat, kernel_init=nn.initializers.glorot_uniform())(xx))
#xx = nn.Dropout(0.3)(xx, deterministic=DETERMINISTIC)
xx = nn.Dense(self.features[-1], use_bias = True)(xx)
return xx
class MLP_constant(nn.Module):
features: Sequence[int]
@nn.compact
def __call__(self, xx):
xx = nn.Dense(self.features[0], use_bias = False)(xx)
return xx
class MLP_nl(nn.Module):
features: Sequence[int]
@nn.compact
def __call__(self, xx):
xx_poly = polynomialFeatures(jnp.reshape(xx, ((-1,))), degree = 3, interaction_only = False, include_bias = False).reshape((-1,)) #jnp.reshape(jnp.power(jnp.reshape(xx, ((-1,1))), jnp.array([1,2,3,4])), dim*4) #
for feat in self.features[:-1]:
xx = nn.swish(nn.Dense(feat, kernel_init=nn.initializers.glorot_uniform(), use_bias = False)(xx))
xx = jnp.concatenate([xx_poly, xx]) #skip connection
xx = nn.Dense(self.features[-1], use_bias = False)(xx)
return xx
class MLP_params(nn.Module):
features: Sequence[int]
sizes_2 = [10,15]
@nn.compact
def __call__(self, xx, in_params):
in_params = jnp.reshape(in_params, ((-1,)))
xx = jnp.concatenate([xx, in_params])
for feat in self.features[:-1]:
xx = nn.tanh(nn.Dense(feat, kernel_init=nn.initializers.glorot_uniform())(xx))
#xx = nn.Dropout(0.3)(xx, deterministic=DETERMINISTIC)
xx = nn.Dense(self.features[-1], use_bias = True)(xx)
#for feat in self.sizes_2:
# in_params = nn.swish(nn.Dense(feat, kernel_init=nn.initializers.glorot_uniform(), use_bias = True)(in_params))
#in_params = nn.Dense(self.features[-1], use_bias = True)(in_params)
return xx
class CNN(nn.Module):
features = [2, 10, 10, 10, 1]
@nn.compact
def __call__(self, xx):
for feat in self.features[:-1]:
xx = nn.Conv(features=feat, kernel_size = (2,2), padding = 'same')(xx)
xx = nn.swish(xx)
xx = nn.Conv(features=self.features[-1], kernel_size = (2,2), padding = 'same')(xx)
return xx
class CNN1D(nn.Module):
features: Sequence[int]
@nn.compact
def __call__(self, xx):
x = xx[0,:LATENT,:]
dx = xx[0,LATENT:,:]
if self.features[0] == 0:
xx = x
xx = jnp.reshape(xx, ((1, LATENT, 1)))
elif self.features[0] == 1:
xx = dx
xx = jnp.reshape(xx, ((1, LATENT, 1)))
xx_poly = jnp.reshape(jnp.power(jnp.reshape(xx, ((-1,1))), jnp.array([1,2])), LATENT*2)
xx = nn.Conv(features=4, kernel_size = (1,2), padding = 'valid', use_bias = False)(xx)
xx = polypool(xx)
xx = jnp.concatenate([xx_poly, xx]) #skip connection
xx = nn.Dense(1, use_bias = True)(xx)
return xx
class MLP_struct(nn.Module):
features: Sequence[int]
@nn.compact
def __call__(self, xx):
x = xx[0,:LATENT,:]
dx = xx[0,LATENT:,:]
term1 = x[0]**2
term2 = x[1]**2
term3 = (x[1] - x[0])**2
term4 = dx[0]**2
term5 = dx[1]**2
term6 = (dx[1] - dx[0])**2
xvec = jnp.array([term1,term2,term3,term4,term5,term6]).reshape((-1,))
xvec = xvec[jnp.array(self.features)]
xx = nn.Dense(1, use_bias = False)(xvec)
return xx
class MLP_poly(nn.Module):
features: Sequence[int]
degree: Sequence[int]
sizes = [6]
@nn.compact
def __call__(self, xx):
x = xx[0,:LATENT,:]
dx = xx[0,LATENT:,:]
if self.features[0] == 0:
xx = x
xx = jnp.reshape(xx, ((LATENT,)))
dim = LATENT
elif self.features[0] == 1:
xx = dx
xx = jnp.reshape(xx, ((LATENT,)))
dim = LATENT
else:
xx = jnp.reshape(xx, ((LATENT*2,)))
dim = LATENT*2
xx_poly = polynomialFeatures(jnp.reshape(xx, ((-1,))), degree = self.degree[0], interaction_only = False, include_bias = False).reshape((-1,)) #jnp.reshape(jnp.power(jnp.reshape(xx, ((-1,1))), jnp.array([1,2,3,4])), dim*4) #
xx_sin = jnp.reshape(jnp.sin(xx), dim)
xx_cos = jnp.reshape(jnp.cos(xx), dim)
for feat in self.sizes[:-1]:
xx = nn.swish(nn.Dense(feat, kernel_init=nn.initializers.glorot_uniform(), use_bias = False)(xx))
xx = nn.swish(nn.Dense(self.sizes[-1], use_bias = False, kernel_init=nn.initializers.glorot_uniform())(xx))
xx = xx_poly #jnp.concatenate([xx_poly, xx]) #skip connection
xx = nn.Dense(1, use_bias = True)(xx)
return xx
class MLP_poly_fom(nn.Module):
features: Sequence[int]
degree: Sequence[int]
sizes = [30,20,10,6]
@nn.compact
def __call__(self, xx):
#x = xx[0,:LATENT,:]
#dx = xx[0,LATENT:,:]
xx = jnp.reshape(xx, ((-1,)))
#xx_poly = polynomialFeatures(jnp.reshape(xx, ((-1,))), degree = self.degree[0], interaction_only = False, include_bias = False).reshape((-1,)) #jnp.reshape(jnp.power(jnp.reshape(xx, ((-1,1))), jnp.array([1,2,3,4])), dim*4) #
#xx_sin = jnp.reshape(jnp.sin(xx), dim)
#xx_cos = jnp.reshape(jnp.cos(xx), dim)
for feat in self.sizes[:-1]:
xx = nn.swish(nn.Dense(feat, kernel_init=nn.initializers.glorot_uniform(), use_bias = False)(xx))
xx = nn.swish(nn.Dense(self.sizes[-1], use_bias = False, kernel_init=nn.initializers.glorot_uniform())(xx))
#xx = xx_poly #jnp.concatenate([xx_poly, xx]) #skip connection
xx = nn.Dense(1, use_bias = True)(xx)
return xx
class MLP_poly_latent(nn.Module):
features: Sequence[int]
nlayers = NLAYERS
nodes = NNODES
sizes = list(nlayers*[nodes])
@nn.compact
def __call__(self, xx):
x = xx[0,:LATENT,:]
dx = xx[0,LATENT:,:]
if self.features[0] == 0:
xx = x
xx = jnp.reshape(xx, ((LATENT,)))
dim = LATENT
elif self.features[0] == 1:
xx = dx
xx = jnp.reshape(xx, ((LATENT,)))
dim = LATENT
else:
xx = jnp.reshape(xx, ((LATENT*2,)))
dim = LATENT*2
xx_poly = jnp.reshape(jnp.power(jnp.reshape(xx, ((-1,1))), jnp.array([1,2,3,4])), dim*4) # polynomialFeatures(jnp.reshape(xx, ((-1,))), degree = 4, interaction_only = False, include_bias = False).reshape((-1,))
xx_sin = jnp.reshape(jnp.sin(xx), dim)
xx_cos = jnp.reshape(jnp.cos(xx), dim)
for feat in self.sizes[:-1]:
xx = nn.swish(nn.Dense(feat, kernel_init=nn.initializers.glorot_uniform(), use_bias = False)(xx))
print(xx.shape, xx_sin.shape, xx_poly.shape)
xx = nn.swish(nn.Dense(self.sizes[-1], use_bias = False, kernel_init=nn.initializers.glorot_uniform())(xx))
xx = jnp.concatenate([xx_poly, xx]) #skip connection
xx = nn.Dense(LATENT, use_bias = True)(xx)
return xx
def hessian(f, argnums):
return jacfwd(jacrev(f, argnums = argnums), argnums = argnums)
def polypool(x):
#x = jnp.sum(x, axis = 1)
#x = jnp.reshape(x, (5,))
x = jnp.power(x, jnp.array([1,2,3,4]))
x = jnp.reshape(x, ((-1,)))
return x
def polynomialFeatures( X, degree = 2, interaction_only = False, include_bias = True ) :
features = X
prev_chunk = X
indices = list( range( len( X ) ) )
for d in range( 1, degree ) :
# Create a new chunk of features for the degree d:
new_chunk = []
# Multiply each component with the products from the previous lower degree:
for i, v in enumerate( X[:-d] if interaction_only else X ) :
# Store the index where to start multiplying with the current component
# at the next degree up:
next_index = len( new_chunk )
for coef in prev_chunk[indices[i+( 1 if interaction_only else 0 )]:] :
new_chunk.append( v*coef )
indices[i] = next_index
# Extend the feature vector with the new chunk of features from the degree d:
features = jnp.append( features, jnp.array(new_chunk))
prev_chunk = new_chunk
if include_bias :
features = jnp.insert( features, 0, 1 )
return features
In [19]:
LATENT = nDof
ORIGINAL_DIM = 64
MSTEPS = 2
METHOD = lm.Adams_Bashforth(MSTEPS)
ALPHA = np.float32(-METHOD.alpha[::-1])
BETA = np.float32(METHOD.beta[::-1])
QMAX = np.max(np.max(QnormTrainData[:,:,:], axis = 0), axis = 0)
XMAX = np.max(np.max(XnormTrainData[:,:,:], axis = 0), axis = 0)
DETERMINISTIC = False
print(XMAX.shape)
MHAT = Mhat
KHAT = Khat
CHAT = Chat
VMAT = jnp.array(V)
DT = dt
class CanonicalTransformer(nn.Module):
def setup(self):
self.E_layers = MLP([10,10,25,LATENT*LATENT])
self.latent = LATENT
self.original = ORIGINAL_DIM
def __call__(self, x, dx, ddx):
x = jnp.reshape(x, (1, self.latent))
dx = jnp.reshape(dx, (1, self.latent))
ddx = jnp.reshape(ddx, (1, self.latent))
in_x = jnp.reshape(x, (self.latent)) #jnp.reshape(jnp.concatenate([x, dx, ddx], axis = 1), (self.original*3))
T_dummy = jnp.reshape(self.E_layers(in_x), (self.latent,self.latent))
T = jnp.eye(self.latent) #jnp.reshape(self.E_layers(in_x), (self.latent,self.latent))
Q = jnp.dot(T,
jnp.reshape(in_x, self.latent))
Q = jnp.reshape(Q, (self.latent,))
return Q
def get_T(self, x):
x = jnp.reshape(x, (1, self.latent))
#dx = jnp.reshape(dx, (1, self.original))
#ddx = jnp.reshape(ddx, (1, self.original))
in_x = jnp.reshape(x, (self.latent))
T = jnp.eye(self.latent)
return T
class InverseCanonicalTransformer(nn.Module):
def setup(self):
self.latent = LATENT
self.original = ORIGINAL_DIM
self.D_layers = MLP_nl([6,6,10,self.original])
def __call__(self, x, dx, ddx):
x = jnp.reshape(x, (1, self.latent, 1))
dx = jnp.reshape(dx, (1, self.latent, 1))
ddx = jnp.reshape(ddx, (1, self.latent, 1))
in_x = jnp.reshape(x, (self.latent)) #jnp.reshape(jnp.concatenate([x, dx, ddx], axis = 1), (self.latent*3))
X_nl= jnp.reshape(self.D_layers(in_x), (self.original,))
T = VMAT
X_lin = jnp.dot(T,
jnp.reshape(in_x, self.latent))
#X_nl = jnp.dot(T_nl,
# jnp.reshape(x_poly, nl_dims))
X = jnp.reshape(X_lin, (self.original,)) + jnp.reshape(X_nl, (self.original,))
return X
def get_T(self, x, dx, ddx):
x = jnp.reshape(x, (1, self.latent, 1))
dx = jnp.reshape(dx, (1, self.latent, 1))
ddx = jnp.reshape(ddx, (1, self.latent, 1))
in_x = jnp.reshape(x, (self.latent))
x_poly = polynomialFeatures(jnp.reshape(x, ((-1,))), degree = 3, interaction_only = False, include_bias = False).reshape((-1,))
x_poly = x_poly[self.latent:]
nl_dims = len(x_poly)
T = jnp.reshape(self.D_layers(in_x), (self.original, nl_dims))
return T
class Kinetic(nn.Module):
def setup(self):
self.T_layers = MLP_poly([1], degree = [2])
self.latent = LATENT
def __call__(self, x, dx):
x = jnp.reshape(x, (1, self.latent, 1))
dx = jnp.reshape(dx, (1, self.latent, 1))
in_x = jnp.concatenate([x, dx], axis = 1)
T = self.T_nl(x, dx)
return T
def T_lin(self, x, dx):
dx = jnp.reshape(dx, (1, self.latent))
T_lin = jnp.dot(dx, jnp.dot(MHAT, dx.T))
return 0.5*jnp.reshape(T_lin, (-1,))
def T_nl(self, x, dx):
in_x = jnp.concatenate([x, dx], axis = 1)
T = jnp.reshape(self.T_layers(in_x), (-1,))
return T
class Potential(nn.Module):
def setup(self):
self.V_layers = MLP_poly_fom(features = [0], degree = [4])
self.latent = LATENT
self.original_dim = ORIGINAL_DIM
def __call__(self, x, dx):
x = jnp.reshape(x, (1, self.original_dim, 1))
dx = jnp.reshape(dx, (1, self.original_dim, 1))
in_x = jnp.concatenate([x, dx], axis = 1)
V = self.T_nl(x, dx)
return V
def T_lin(self, x, dx):
x = jnp.reshape(x, (1, self.latent))
T_lin = 0.5*jnp.dot(x, jnp.dot(KHAT, x.T))
return jnp.reshape(T_lin, (-1,))
def T_nl(self, x, dx):
in_x = jnp.concatenate([x, dx], axis = 1)
T = jnp.reshape(self.V_layers(x), (-1,)) # only pass displacements
return T
class Rayleigh(nn.Module):
def setup(self):
self.R_layers = MLP_poly([1], degree = [2])
self.latent = LATENT
def __call__(self, x, dx):
x = jnp.reshape(x, (1, self.latent, 1))
dx = jnp.reshape(dx, (1, self.latent, 1))
in_x = jnp.concatenate([x, dx], axis = 1)
R = self.T_nl(x, dx)
return R
def T_lin(self, x, dx):
dx = jnp.reshape(dx, (1, self.latent))
T_lin = 0.5*jnp.dot(dx, jnp.dot(CHAT, dx.T))
return jnp.reshape(T_lin, (-1,))
def T_nl(self, x, dx):
in_x = jnp.concatenate([x, dx], axis = 1)
T = jnp.reshape(self.R_layers(in_x), (-1,))
return T
class Qnn(nn.Module):
def setup(self):
self.Qnn_layers = MLP_poly_latent([2])
self.latent = LATENT
def __call__(self, x, dx):
x = jnp.reshape(x, (1, self.latent, 1))
dx = jnp.reshape(dx, (1, self.latent, 1))
in_x = jnp.concatenate([x, dx], axis = 1)
V = jnp.reshape(self.Qnn_layers(in_x), (-1,))
return V*0.0
class Minv_CNN(nn.Module):
def setup(self):
self.cnn_layers = CNN()
self.latent = LATENT
def __call__(self, Mmat):
Mmat_re = jnp.reshape(Mmat, (self.latent, self.latent, 1))
Minv_ = jnp.reshape(self.cnn_layers(Mmat_re), (self.latent, self.latent))
return Minv_
class AEROM(nn.Module):
def setup(self):
self.T_layers = Kinetic()
self.V_layers = Potential()
self.R_layers = Rayleigh()
self.Qnn_layers = Qnn()
self.encoderNet_q = CanonicalTransformer()
#self.encoderNet_dq = CanonicalTransformer()
#self.encoderNet_ddq = CanonicalTransformer()
self.decoderNet_q = InverseCanonicalTransformer()
#self.decoderNet_dq = InverseCanonicalTransformer()
#self.decoderNet_ddq = InverseCanonicalTransformer()
self.latent = LATENT
def __call__(self, x, dx, ddx):
#qdot = self.encoderNet_dq(x, dx, ddx)
#qdotdot = self.encoderNet_ddq(x, dx, ddx)
T = self.T_layers(x,x)
R = self.R_layers(x,x)
q_ = self.encoderNet_q(x, x, x) #dummy call
x_ = self.decoderNet_q(x, x, x) #dummy call
V = self.V_layers(x_,x_)
Qnn = self.Qnn_layers(x,x)
#dx_ = self.decoderNet_dq(q, qdot, qdotdot)
#ddx_ = self.decoderNet_ddq(q, qdot, qdotdot)
return T
def get_quadratic_energy_jnp(K,x):
xi = x.reshape((1,LATENT))
#print(K.shape, xi.shape)
E = 1/2.*jnp.dot(xi, jnp.dot(K, xi.T))
return jnp.reshape(E, (-1,))
def Lagrangian(params, x, dx):
T = eval_kinetic(params, x, dx) #Kinetic().apply({'params': params[0]}, x, dx)
V = eval_potential(params, x, dx) #Potential().apply({'params': params[1]}, x, dx)
return T-V
def eval_kinetic(params,x, dx):
T = Kinetic().apply({'params': params[0]}, x, dx)
B = CanonicalTransformer().apply({'params': params[4]}, x, method = CanonicalTransformer().get_T)
A = jnp.linalg.pinv(jnp.reshape(B, (LATENT,LATENT)))
Mtrans = jnp.dot(jnp.transpose(A), jnp.dot(MHAT, A))
print(B.shape, A.shape, MHAT.shape, dx.shape, Mtrans.shape)
T_lin = 0.5*jnp.dot(dx, jnp.dot(Mtrans, dx.T))
return T*0.0 + T_lin
def eval_potential(params, x, dx):
xfom = jnp.dot(VMAT, jnp.reshape(x, (LATENT, 1)))
dxfom = jnp.dot(VMAT, jnp.reshape(dx, (LATENT, 1)))
V = Potential().apply({'params': params[1]}, xfom, dxfom)
B = CanonicalTransformer().apply({'params': params[4]}, x, method = CanonicalTransformer().get_T)
A = jnp.linalg.pinv(jnp.reshape(B, (LATENT,LATENT)))
Ktrans = jnp.dot(jnp.transpose(A), jnp.dot(KHAT, A))
V_lin = 0.5*jnp.dot(x, jnp.dot(Ktrans, x.T))
return V + V_lin
def eval_rayleigh(params, x, dx):
R = Rayleigh().apply({'params': params[2]}, x, dx)
B = CanonicalTransformer().apply({'params': params[4]}, x, method = CanonicalTransformer().get_T)
A = jnp.linalg.pinv(jnp.reshape(B, (LATENT,LATENT)))
Ctrans = jnp.dot(jnp.transpose(A), jnp.dot(CHAT, A))
R_lin = 0.5*jnp.dot(dx, jnp.dot(Ctrans, dx.T))
return (R + R_lin)*0.0
def eval_kinetic_lin(params,x, dx):
B = CanonicalTransformer().apply({'params': params[4]}, x, method = CanonicalTransformer().get_T)
A = jnp.linalg.pinv(jnp.reshape(B, (LATENT,LATENT)))
Mtrans = jnp.dot(jnp.transpose(A), jnp.dot(MHAT, A))
T_lin = 0.5*jnp.dot(dx, jnp.dot(Mtrans, dx.T))
return T_lin
def eval_potential_lin(params, x, dx):
B = CanonicalTransformer().apply({'params': params[4]}, x, method = CanonicalTransformer().get_T)
A = jnp.linalg.pinv(jnp.reshape(B, (LATENT,LATENT)))
Ktrans = jnp.dot(jnp.transpose(A), jnp.dot(KHAT, A))
V_lin = 0.5*jnp.dot(x, jnp.dot(Ktrans, x.T))
return V_lin
def eval_rayleigh_lin(params, x, dx):
B = CanonicalTransformer().apply({'params': params[4]}, x, method = CanonicalTransformer().get_T)
A = jnp.linalg.pinv(jnp.reshape(B, (LATENT,LATENT)))
Ctrans = jnp.dot(jnp.transpose(A), jnp.dot(CHAT, A))
R_lin = 0.5*jnp.dot(dx, jnp.dot(Ctrans, dx.T))
return R_lin
def eval_kinetic_nl(params,x, dx):
T = Kinetic().apply({'params': params[0]}, x, dx)
return T
def eval_potential_nl(params, x, dx):
xfom = jnp.dot(VMAT, jnp.reshape(x, (LATENT, 1)))
dxfom = jnp.dot(VMAT, jnp.reshape(dx, (LATENT, 1)))
V = Potential().apply({'params': params[1]}, xfom, dxfom)
return V
def eval_rayleigh_nl(params, x, dx):
R = Rayleigh().apply({'params': params[2]}, x, dx)
return R
def encode_q(E_params, x, dx, ddx):
R = CanonicalTransformer().apply({'params': E_params[0]}, x, dx, ddx)
return R
def eval_Qnn(params, x, dx):
R = Qnn().apply({'params': params[3]}, x, dx)
return R
def encode_dq(E_params, x, dx, ddx):
dx_dq = jax.jacfwd(encode_q, argnums = 1)(E_params, x, dx, ddx) #dx/dq
dx = jnp.reshape(dx, (LATENT,))
dX = jnp.matmul(dx_dq, dx)
return jnp.reshape(dX, (LATENT,))
def encode_ddq(E_params, x, dx, ddx):
ddx_ddq = jax.jacfwd(encode_dq, argnums = 2)(E_params, x, dx, ddx) #dxdot/dqdot
ddx_dq = jax.jacfwd(encode_dq, argnums = 1)(E_params, x, dx, ddx) #dxdot/dq
ddx = jnp.reshape(ddx, (LATENT,))
dx = jnp.reshape(dx, (LATENT,))
ddX = jnp.matmul(ddx_ddq, ddx) + jnp.matmul(ddx_dq, dx)
return jnp.reshape(ddX, (LATENT,))
def decode_q(D_params, x, dx, ddx):
R = InverseCanonicalTransformer().apply({'params': D_params[0]}, x, dx, ddx)
return R
def get_T_x(E_params, x, dx, ddx, comp=0):
T = InverseCanonicalTransformer().apply({'params': E_params[comp]}, x, dx, ddx, method = InverseCanonicalTransformer().get_T)
return T
def decode_dq(D_params, x, dx, ddx):
dx_dq = jax.jacfwd(decode_q, argnums = 1)(D_params, x, dx, ddx) #dx/dq
dx = jnp.reshape(dx, (LATENT,))
dX = jnp.matmul(dx_dq, dx)
return jnp.reshape(dX, (ORIGINAL_DIM,))
def decode_ddq(D_params, x, dx, ddx):
ddx_ddq = jax.jacfwd(decode_dq, argnums = 2)(D_params, x, dx, ddx) #dxdot/dqdot
ddx_dq = jax.jacfwd(decode_dq, argnums = 1)(D_params, x, dx, ddx) #dxdot/dq
ddx = jnp.reshape(ddx, (LATENT,))
dx = jnp.reshape(dx, (LATENT,))
ddX = jnp.matmul(ddx_ddq, ddx) + jnp.matmul(ddx_dq, dx)
return jnp.reshape(ddX, (ORIGINAL_DIM,))
def encoder(E_params, x, dx, ddx):
q = encode_q(E_params, x, dx, ddx)
dq = encode_dq(E_params, x, dx, ddx)
ddq = encode_ddq(E_params, x, dx, ddx)
return q, dq, ddq
def decoder(D_params, x, dx, ddx):
q = decode_q(D_params, x, dx, ddx)
dq = decode_dq(D_params, x, dx, ddx)
ddq = decode_ddq(D_params, x, dx, ddx)
return q, dq, ddq
def init_params(data):
rng = random.PRNGKey(0)
rng, key = random.split(rng)
pt = Kinetic().init(key, data[:1,:,0], data[:1,:,1])['params']
pv = Potential().init(key, data[:1,:,0], data[:1,:,1])['params']
pM = Minv_CNN().init(key, jnp.zeros((1, LATENT, LATENT, 1)))['params']
return [pt, pv, pM]
def momentum(params, x, dx):
latent = LATENT
dx = jnp.reshape(dx, (-1, latent))
x = jnp.reshape(x, (-1, latent))
p = jax.jacfwd(Lagrangian, argnums = 2)(params, x, dx)
p = jnp.reshape(p, (latent,))
return p
def Hamiltonian(params, q, qdot, p):
latent = LATENT
qdot = jnp.reshape(qdot, (-1, latent))
q = jnp.reshape(q, (-1, latent))
L = Lagrangian(params,q,qdot)
qdotp = jnp.dot(qdot, jnp.reshape(p, (latent, -1)))
qdotp = jnp.reshape(qdotp, (-1,))
H = qdotp - L
return H
def canonical(params, q, qdot, p, Qext):
latent = LATENT
dH_dp = jax.jacfwd(Hamiltonian, argnums = 3)(params, q, qdot, p)
dH_dq = jax.jacfwd(Hamiltonian, argnums = 1)(params, q, qdot, p)
dR_dqdot = jax.jacfwd(eval_rayleigh, argnums = 2)(params, q, qdot)
Qext = jnp.reshape(Qext, ((latent,)))
dq_dt = (dH_dp).reshape((latent,))
dp_dt = (-dH_dq - dR_dqdot + Qext).reshape((latent,))
return [dq_dt, dp_dt]
def EulerLagrange(params, q, qdot, Qext):
latent = LATENT
p = momentum(params, q, qdot)
[_, dp] = canonical(params, q, qdot, p, Qext)
dL_dqdot = jax.jacfwd(Lagrangian, argnums = 2)(params, q, qdot)
dL_dq = jax.jacfwd(Lagrangian, argnums = 1)(params, q, qdot)
dR_dqdot = jax.jacfwd(eval_rayleigh, argnums = 2)(params, q, qdot)
d_dL_dqdot_dt = dp #(tape.gradient(L, q) - tape.gradient(R, qdot) + Fext)
EL_resid = jnp.reshape(d_dL_dqdot_dt, (latent,)) - jnp.reshape(dL_dq, (latent,)) + jnp.reshape(dR_dqdot, (latent,)) - jnp.reshape(Qext, (latent,))
#dt_check = F_loss(dL_dqdot, d_dL_dqdot_dt, dt = DT)
return jnp.mean(jnp.square(EL_resid)) #+ dt_check*1.0e-3
def get_Mass_Matrix(params, q, dq):
latent = LATENT
M = jax.hessian(eval_kinetic, argnums = 2)(params, q, dq)
M = jnp.reshape(M, (latent, latent))
return M #jnp.reshape(M, (self.latent, self.latent))
def get_K_Matrix(params, q, dq):
latent = LATENT
M = jax.hessian(eval_potential, argnums = 1)(params, q, dq)
M = jnp.reshape(M, (latent, latent))
return M #jnp.reshape(M, (self.latent, self.latent))
def get_C_Matrix(params, q, dq):
latent = LATENT
M = jax.hessian(eval_rayleigh, argnums = 2)(params, q, dq)
M = jnp.reshape(M, (latent, latent))
return M #jnp.reshape(M, (self.latent, self.latent))
def get_Mass_Matrix_nl(params, q, dq):
latent = LATENT
M = jax.hessian(eval_kinetic_nl, argnums = 2)(params, q, dq)
M = jnp.reshape(M, (latent, latent))
return M #jnp.reshape(M, (self.latent, self.latent))
def get_K_Matrix_nl(params, q, dq):
latent = LATENT
M = jax.hessian(eval_potential_nl, argnums = 1)(params, q, dq)
M = jnp.reshape(M, (latent, latent))
return M #jnp.reshape(M, (self.latent, self.latent))
def get_C_Matrix_nl(params, q, dq):
latent = LATENT
M = jax.hessian(eval_rayleigh_nl, argnums = 2)(params, q, dq)
M = jnp.reshape(M, (latent, latent))
return M #jnp.reshape(M, (self.latent, self.latent))
def get_dM_dq(params, q, dq):
dM = jax.jacfwd(get_Mass_Matrix, argnums = 1)(params, q, dq)
return dM
def get_Minv(params, Mmat):
latent = LATENT
Minv = jnp.linalg.pinv(jnp.reshape(Mmat, (latent, latent)))
return Minv
def get_Qnc(params, q, dq, Qext, in_params):
Qnn_ = eval_Qnn(params, q, dq, in_params)
# Qnn_vec = jnp.zeros(Qext.shape)
#x.at[idx].set(y) syntax
#Qnn_vec.at[NDZ].set(Qnn_)
Qnc = Qext - Qnn_
return Qnc
def get_Qnn(params, q, dq, in_params):
Qnn_ = eval_Qnn(params, q, dq, in_params)
return Qnn_
def get_Qforce(D_params, q, dq, ddq, Fext):
dx_dq = jax.jacfwd(decode_q, argnums = 1)(D_params, q, dq, ddq)
dx_dq = jnp.transpose(dx_dq, (1,0))
Fext = jnp.reshape(Fext, (ORIGINAL_DIM,))
Qext = jnp.matmul(dx_dq, Fext)
return jnp.reshape(Qext, (LATENT,))
def get_ddq(params, q, dq, dp):
latent = LATENT
Mmat = get_Mass_Matrix(params, q, dq)
dM = get_dM_dq(params, q, dq)
Mdq = jnp.matmul(dM, jnp.reshape(dq, (1,latent,1)))
term2 = jnp.matmul(jnp.reshape(Mdq, (latent,latent)),
jnp.reshape(dq, (latent,1)))
term1 = jnp.reshape(dp, (latent,))
term2 = jnp.reshape(term2, (latent,))
Minv = get_Minv(params, Mmat)
ddq = jnp.matmul(Minv, jnp.reshape(term1 - term2, (latent, 1)))
ddq = jnp.reshape(ddq, (latent,))
return ddq
def general_checks(params, q, dq):
latent = LATENT
Mmat = get_Mass_Matrix(params, q, dq)
phat = jnp.matmul(jnp.reshape(Mmat, (latent, latent)),
jnp.reshape(dq, (latent, 1))
)
p = momentum(params, q, dq)
p_check = jnp.square(jnp.reshape(p, (latent,)) - jnp.reshape(phat, (latent,)))
##
qdotp = jnp.dot(jnp.reshape(dq, (1, latent)),
jnp.reshape(p, (latent, 1))
)
That = 1/2.*qdotp
T = eval_kinetic(params, q, dq)
V = eval_potential(params, q, dq)
T_check = jnp.square(jnp.reshape(T, (-1,)) - jnp.reshape(That, (-1,)))
##
H = Hamiltonian(params, q, dq, p)
Hhat = T + V
H_check = jnp.square(H - Hhat)
##
#Minv = get_Minv(params, Mmat)
#eye_hat = jnp.matmul(jnp.reshape(Mmat, (latent, latent)),
# jnp.reshape(Minv, (latent, latent))
#)
#eye = jnp.eye(latent)
#Minv_check = jnp.sum(jnp.square(
# eye_hat - eye
#))
##
Veq = eval_potential(params, jnp.zeros(q.shape), jnp.zeros(dq.shape))
Veq_check = jnp.square(Veq)
##
neg_M = jnp.mean(jnp.square(jax.nn.relu(jnp.negative(Mmat))))
neg_T = jnp.mean(jnp.square(jax.nn.relu(jnp.negative(T))))
neg_V = jnp.mean(jnp.square(jax.nn.relu(jnp.negative(V))))
neg_checks = neg_M + neg_T + neg_V
Mnorm = jnp.linalg.norm(Mmat)
Mnorm_loss = jnp.square(Mnorm.reshape((1,)) - np.linalg.norm(np.array([[1.1,0],[0,1.1]]))/Mnorm)
##
ix_off = ~np.eye(Mmat.shape[0],dtype=bool)
offM = jnp.sum(jnp.square(Mmat[ix_off]))
##
M11 = jnp.square(Mmat[0,0] - 1.0)
##
#[-0.1,0.5, 0.1, 0]
#x0 = jnp.array(w0[:LATENT])
#dx0 = jnp.array(w0[LATENT:])
#L0 = Lagrangian(params, x0, dx0)
#L0_check = jnp.mean(jnp.square(L0 - L0_analyt))
return jnp.mean(T_check)*10 + jnp.mean(p_check)*10 + jnp.mean(H_check)*10 #+ L0_check*0.1 #+ neg_checks*1 #+ H_check*0.1 #+ jnp.mean(M11)# + offM #+ Mnorm_loss
def F_loss(p, dp, dt):
alpha = ALPHA
beta = BETA
M = MSTEPS
#dt = DT
p = jnp.reshape(p, (-1, LATENT))
dp = jnp.reshape(dp, (-1, LATENT))
pmax = jnp.max(jnp.array([jnp.max(jnp.abs(p)), 1])) # trying to normalize so amplitudes are better
p = p/pmax
dp = dp/pmax
Y = alpha[0]*p[M:, :] + dt*beta[0]*dp[M:, :]
for m in range(1, M+1):
Y = Y + alpha[m]*p[M-m:-m, :] + dt*beta[m]*dp[M-m:-m,:]
return jnp.mean(jnp.square(Y))
def jnp_cov(x):
nDof = LATENT
mean_x = jnp.reshape(jnp.mean(x, axis = 0), (1, nDof))
mx = jnp.matmul(mean_x.T, mean_x)
vx = jnp.matmul(x.T, x) #/tf.cast(tf.shape(x)[0], tf.float32)
cov_xx = vx - mx
return cov_xx
In [20]:
def l2_loss(x, alpha):
return alpha * (x ** 2).mean()
def l1_loss(x, alpha):
return alpha * jnp.abs(x).sum()
@jax.jit
def f_diff(params, x, dx, Fext):
latent = LATENT
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params, params['encoderNet_q']]
#Qnc = jax.vmap(get_Qnc, in_axes = (None, 0, 0, 0))(params_, x, dx, Fext, in_params)
p = jax.vmap(momentum, in_axes = (None, 0, 0))(params_, x, dx)
[dq,dp] = jax.vmap(canonical, in_axes = (None, 0, 0, 0, 0))(params_, x, dx, p, Fext)
ddq = jax.vmap(get_ddq, in_axes = (None, 0, 0, 0))(params_, x, dx, dp)
X = jnp.concatenate([jnp.reshape(dx, (-1, latent)),
jnp.reshape(ddq, (-1, latent))], axis = -1)
return X
def check_1(params, batch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
params_ = [T_params, V_params, R_params]
q = batch[:, :, 0]
qdot = batch[:, :, 1]
ddq = batch[:, :, 2]
Qext = batch[:, :, 3]
p = jax.vmap(momentum, in_axes = (None, 0, 0))(params_, q, qdot)
[dq, dp] = jax.vmap(canonical, in_axes = (None, 0, 0, 0, 0))(params_, q, qdot, p, Qext)
dL_dqdot = jax.vmap(jax.jacfwd(Lagrangian, argnums = 2), in_axes = (None, 0, 0))(params_, q, qdot)
dL_dq = jax.vmap(jax.jacfwd(Lagrangian, argnums = 1), in_axes = (None, 0, 0))(params_, q, qdot)
d_dL_dqdot_dt = dp #(tape.gradient(L, q) - tape.gradient(R, qdot) + Fext)
return [p,dp,dL_dqdot,dL_dq,dq]
def encode_trans(params, batch, comp=0):
E_params = [params['encoderNet_q']]
D_params = [params['decoderNet_q']]
X = batch[:, :, 0]
dX = batch[:, :, 1]
ddX = batch[:, :, 2]
Fext = batch[:, :, 3]
T = jax.vmap(get_T_q, in_axes = (None, 0, 0, 0, None))(E_params, X, dX, ddX, comp)
return T
def decode_trans(params, batch,comp=0):
D_params = [params['decoderNet_q']]
Q = batch[:, :, 0]
dQ = batch[:, :, 1]
ddQ = batch[:, :, 2]
T = jax.vmap(get_T_x, in_axes = (None, 0, 0, 0, None))(D_params, Q, dQ, ddQ, comp)
return T
def encode_predict(params, batch):
E_params = [params['encoderNet_q']]
D_params = [params['decoderNet_q']]
X = batch[:, :, 0]
dX = batch[:, :, 1]
ddX = batch[:, :, 2]
Fext = batch[:, :, 3]
x,dx,ddx = jax.vmap(encoder, in_axes = (None, 0, 0, 0))(E_params, X, dX, ddX)
#Qext = jax.vmap(get_Qforce, in_axes = (None, 0, 0, 0, 0))(D_params, x, dx, ddx, Fext)
Q = jnp.concatenate([jnp.reshape(x, (-1,LATENT,1)),
jnp.reshape(dx, (-1,LATENT,1)),
jnp.reshape(ddx, (-1,LATENT,1)),
jnp.reshape(Fext, (-1,LATENT,1))], axis = -1)
return Q
def decode_predict(params, batch):
D_params = [params['decoderNet_q']]
Q = batch[:, :, 0]
dQ = batch[:, :, 1]
ddQ = batch[:, :, 2]
x,dx,ddx = jax.vmap(decoder, in_axes = (None, 0, 0, 0))(D_params, Q, dQ, ddQ)
X = jnp.concatenate([jnp.reshape(x, (-1,ORIGINAL_DIM,1)),
jnp.reshape(dx, (-1,ORIGINAL_DIM,1)),
jnp.reshape(ddx, (-1,ORIGINAL_DIM,1))], axis = -1)
return X
@jax.jit
def predict_ddq(params, batch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params, params['encoderNet_q']]
x = batch[:, :, 0]
dx = batch[:, :, 1]
ddx = batch[:, :, 2]
Qext = batch[:, :, 3]
#Qnc = jax.vmap(get_Qnc, in_axes = (None, 0, 0, 0, 0))(params_, x, dx, Qext, in_params)
p = jax.vmap(momentum, in_axes = (None, 0, 0))(params_, x, dx)
[_,dp] = jax.vmap(canonical, in_axes = (None, 0, 0, 0, 0))(params_, x, dx, p, Qext)
ddq = jax.vmap(get_ddq, in_axes = (None, 0, 0, 0))(params_, x, dx, dp)
return ddq
def predict_Qnc(params, batch, pbatch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params]
x = batch[:, :, 0]
dx = batch[:, :, 1]
ddx = batch[:, :, 2]
Qext = batch[:, :, 3]
in_params = pbatch[:,:]
Qnc = jax.vmap(get_Qnc, in_axes = (None, 0, 0, 0, 0))(params_, x, dx, Qext, in_params)
return Qnc
def predict_kinetic(params, batch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params, params['encoderNet_q']]
x = batch[:, :, 0]
dx = batch[:, :, 1]
ddx = batch[:, :, 2]
T = jax.vmap(eval_kinetic, in_axes = (None, 0, 0))(params_, x, dx)
T_lin = jax.vmap(eval_kinetic_lin, in_axes = (None, 0, 0))(params_, x, dx)
T_nl = jax.vmap(eval_kinetic_nl, in_axes = (None, 0, 0))(params_, x, dx)
return T,T_lin,T_nl
def predict_potential(params, batch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params, params['encoderNet_q']]
x = batch[:, :, 0]
dx = batch[:, :, 1]
ddx = batch[:, :, 2]
V = jax.vmap(eval_potential, in_axes = (None, 0, 0))(params_, x, dx)
V_lin = jax.vmap(eval_potential_lin, in_axes = (None, 0, 0))(params_, x, dx)
V_nl = jax.vmap(eval_potential_nl, in_axes = (None, 0, 0))(params_, x, dx)
return V,V_lin,V_nl
def predict_Hamiltonian(params, batch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params, params['encoderNet_q']]
x = batch[:, :, 0]
dx = batch[:, :, 1]
ddx = batch[:, :, 2]
p = jax.vmap(momentum, in_axes = (None, 0, 0))(params_, x, dx)
H = jax.vmap(Hamiltonian, in_axes = (None, 0, 0, 0))(params_, x, dx, p)
return H
def predict_mass(params, batch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params, params['encoderNet_q']]
x = batch[:, :, 0]
dx = batch[:, :, 1]
ddx = batch[:, :, 2]
M = jax.vmap(get_Mass_Matrix, in_axes = (None, 0, 0))(params_, x, dx)
Minv = jax.vmap(get_Minv, in_axes = (None, 0))(params_, M)
return M, Minv
def predict_MKC(params, batch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params, params['encoderNet_q']]
x = batch[:, :, 0]
dx = batch[:, :, 1]
ddx = batch[:, :, 2]
M = jax.vmap(get_Mass_Matrix, in_axes = (None, 0, 0))(params_, x, dx)
K = jax.vmap(get_K_Matrix, in_axes = (None, 0, 0))(params_, x, dx)
C = jax.vmap(get_C_Matrix, in_axes = (None, 0, 0))(params_, x, dx)
M_nl = jax.vmap(get_Mass_Matrix_nl, in_axes = (None, 0, 0))(params_, x, dx)
K_nl = jax.vmap(get_K_Matrix_nl, in_axes = (None, 0, 0))(params_, x, dx)
C_nl = jax.vmap(get_C_Matrix_nl, in_axes = (None, 0, 0))(params_, x, dx)
return M,K,C,M_nl,K_nl,C_nl
def compute_metrics(loss):
metrics = {
'loss': loss
}
return metrics
def lasso_loss(params, alpha = 1e-3):
T_params = params['T_layers']
V_params = params['V_layers']
wT = Kinetic().apply({'params': T_params}, method = Kinetic().get_coeffs)
wV = Potential().apply({'params': V_params}, method = Potential().get_coeffs)
return l1_loss(wT, alpha) + l1_loss(wV, alpha)
def loss_fn_(params, batch, xbatch):
T_params = params['T_layers']
V_params = params['V_layers']
R_params = params['R_layers']
Qnn_params = params['Qnn_layers']
params_ = [T_params, V_params, R_params, Qnn_params, params['encoderNet_q']]
E_params = [params['encoderNet_q']]
D_params = [params['decoderNet_q']]
x_ = batch[:, :, 0]
dx_ = batch[:, :, 1]
ddx_ = batch[:, :, 2]
Qext = batch[:, :, 3]
X = xbatch[:, :, 0]
dX = xbatch[:, :, 1]
ddX = xbatch[:, :, 2]
print(x_.shape, X.shape)
x,dx,ddx = jax.vmap(encoder, in_axes = (None, 0, 0, 0))(E_params, x_,dx_,ddx_)
X_, dX_, ddX_ =jax.vmap(decoder, in_axes = (None, 0, 0, 0))(D_params, x, dx, ddx)
T_nl = jax.tree_util.tree_leaves(params['decoderNet_q'])[-1].T
T_lin = VMAT
tangency = jnp.mean(jnp.dot(T_lin.T, T_nl)**2)
print(x.shape, dx.shape)
p = jax.vmap(momentum, in_axes = (None, 0, 0))(params_, x, dx)
H = jax.vmap(Hamiltonian, in_axes = (None, 0, 0, 0))(params_, x, dx, p)
H_std = jnp.mean(jnp.diff(H, axis = 0))
[dq,dp] = jax.vmap(canonical, in_axes = (None, 0, 0, 0, 0))(params_, x, dx, p, Qext)
#EL_loss = jax.vmap(EulerLagrange, in_axes = (None, 0, 0, 0))(params_, x, dq, Qext)
ddq = jax.vmap(get_ddq, in_axes = (None, 0, 0, 0))(params_, x, dq, dp)
dL_dqdot = jax.vmap(jax.jacfwd(Lagrangian, argnums = 2), in_axes = (None, 0, 0))(params_, x, dq)
F_loss_dp = F_loss(p, dp, dt = DT)*1000
F_loss_dx = F_loss(dq, ddq, dt = DT)*1000
F_loss_x = F_loss(x, dq, dt = DT)*1000
recon_loss_q = jnp.mean(jnp.square((X[:,:] - X_[:,:])/XMAX[0]))
recon_loss_dq = jnp.mean(jnp.square((dX[:,:] - dX_[:,:])/XMAX[1]))
recon_loss_ddq = jnp.mean(jnp.square((ddX[:,:] - ddX_[:,:])/XMAX[2]))
recon_loss_ddq_enc = jnp.mean(jnp.square((ddx - ddq)/QMAX[2]))
recon_loss = (recon_loss_ddq + recon_loss_dq + recon_loss_q)
checks_loss = jax.vmap(general_checks, in_axes = (None, 0, 0))(params_, x, dq)
d_dL_dqdot_dt = dp #(tape.gradient(L, q) - tape.gradient(R, qdot) + Fext)
EL_check = F_loss(dL_dqdot, d_dL_dqdot_dt, dt = DT)
#supervQ = jnp.mean(jnp.square((X - x)/XMAX[0]))*0.0 #+ jnp.mean(jnp.square((X[:,3] - x[:,2])/XMAX[0])) #force the first latent dim to be the same
loss = jnp.array([F_loss_dp, F_loss_x, F_loss_dx, jnp.mean(checks_loss)*5, jnp.mean(EL_check), recon_loss_ddq_enc*800, 0.1*recon_loss, 0.1*tangency, H_std])
return loss
def loss_fn(params, batch, xbatch):
loss = loss_fn_(params, batch, xbatch)
loss = jnp.sum(loss)
loss += sum(
l2_loss(w, alpha=5e-3)
for w in jax.tree_leaves(params)
)
return loss
@jax.jit
def train_step(state, batch, xbatch):
grad_fn = jax.value_and_grad(loss_fn, has_aux=False)
loss, grads = grad_fn(state.params, batch, xbatch)
state = state.apply_gradients(grads=grads)
metrics = compute_metrics(loss)
return state, metrics
def train_epoch(state, train_ds, xtrain_ds, batch_size, epoch, rng):
"""Train for a single epoch."""
train_ds_size = len(train_ds)
steps_per_epoch = train_ds_size // batch_size
train_ds_batched = train_ds.batch(batch_size).prefetch(1)
train_ds_batched = tfds.as_numpy(train_ds_batched)
xtrain_ds_batched = xtrain_ds.batch(batch_size).prefetch(1)
xtrain_ds_batched = tfds.as_numpy(xtrain_ds_batched)
batch_metrics = []
for batch, xbatch in zip(train_ds_batched, xtrain_ds_batched):
state, metrics = train_step(state, batch, xbatch)
batch_metrics.append(metrics)
# compute mean of metrics across each batch in epoch.
batch_metrics_np = jax.device_get(batch_metrics)
epoch_metrics_np = {
k: np.mean([metrics[k] for metrics in batch_metrics_np])
for k in batch_metrics_np[0]}
if epoch % 100 == 0:
print('train epoch: %d, loss: %.8f' % (
epoch, epoch_metrics_np['loss']))
return state
@jax.jit
def eval_step(params, batch, xbatch):
loss = loss_fn(params, batch, xbatch)
return compute_metrics(loss)
def eval_model(params, test_ds, xtest_ds):
test_ds_batched = test_ds.batch(len(test_ds)).prefetch(1)
test_ds_batched = tfds.as_numpy(test_ds_batched)
xtest_ds_batched = xtest_ds.batch(len(xtest_ds)).prefetch(1)
xtest_ds_batched = tfds.as_numpy(xtest_ds_batched)
for batch, xbatch in zip(test_ds_batched, xtest_ds_batched):
metrics = eval_step(params, batch, xbatch)
metrics = jax.device_get(metrics)
summary = jax.tree_map(lambda x: x.item(), metrics)
return summary['loss']
In [21]:
rng = random.PRNGKey(0)
rng, key = random.split(rng)
init_rng = {'params': random.PRNGKey(0), 'dropout': random.PRNGKey(1)}
initial_variables = AEROM().init(rng, QnormTrainData[:1,:,0], QnormTrainData[:1,:,1], QnormTrainData[:1,:,0])
state = train_state.TrainState.create(
apply_fn = AEROM().apply,
params = initial_variables['params'],
tx = optax.adam(
learning_rate=1e-4,
b1=0.9,
b2=0.98,
eps=1e-7)
)
In [22]:
#jax.tree_util.tree_leaves(initial_variables['params']['decoderNet_q'])[-1].shape
In [52]:
from flax.training import checkpoints
dir_name = "./chkpts/chkpts_DQuinn_NOQnn_v01_AE_decTenc_UPDdof6_DecIdent_v03_PolyOnly_noDamp_NLManifold_Vq_diffH_upd/"
load = 0
purge = 1
if load:
test = os.listdir(dir_name)
chk_latest = checkpoints.latest_checkpoint(dir_name)
state = checkpoints.restore_checkpoint(chk_latest, state)
if purge:
test = os.listdir(dir_name)
for item in test:
os.remove(os.path.join(dir_name, item))
In [53]:
from flax.training import checkpoints
num_epochs = 150000
batch_size = 250
for epoch in range(1, num_epochs + 1):
# Use a separate PRNG key to permute image data during shuffling
rng, input_rng = jax.random.split(rng)
# Run an optimization step over a training batch
state = train_epoch(state, q_train_dataset, x_train_dataset, batch_size, epoch, input_rng)
# Evaluate on the test set after each training epoch
test_loss= eval_model(state.params, q_test_dataset, x_test_dataset)
if epoch % 50 == 0:
print(' test epoch: %d, loss: %.8f' % (
epoch, test_loss))
if epoch % 50 == 0:
checkpoints.save_checkpoint(dir_name, state, epoch, keep=3)
In [56]:
Nsim = 35000
Qpred = encode_predict(state.params, QnormTrainData[:Nsim,:,:])
Xpred = decode_predict(state.params, Qpred)
print(Qpred.shape)
sample = np.array([0,10,30])
plt.figure(figsize = (16,12))
plt.subplot(3,2,1)
plt.plot(Qpred[:,:,0])
plt.plot(QnormTrainData[:Nsim,:,0], 'k--')
#plt.plot(normTrainData[:Nsim,IX_NODELAY[0],0], 'k--')
#plt.plot(QnormTrainData[:Nsim,:,0], 'm--')
plt.subplot(3,2,2)
plt.plot(Xpred[:,sample,0])
plt.plot(XnormTrainData[:Nsim,sample,0], 'k--')
plt.subplot(3,2,3)
plt.plot(Qpred[:,:,1])
#plt.plot(QnormTrainData[:1000,:,1], 'k--')
#dQpred = np.diff(Qpred[:,:,0], axis = 0)/dt
#plt.plot(dQpred, 'm--')
plt.subplot(3,2,4)
plt.plot(Xpred[:,sample,1])
plt.plot(XnormTrainData[:Nsim,sample,1], 'k--')
plt.subplot(3,2,5)
plt.plot(Qpred[:,:,2])
#plt.plot(QnormTrainData[:1000,:,2], 'k--')
plt.subplot(3,2,6)
plt.plot(Xpred[:,sample,2])
plt.plot(XnormTrainData[:Nsim,sample,2], 'k--')
plt.show()
In [57]:
loss_fn_(state.params, QnormTrainData[:Nsim,:,:], XnormTrainData[:Nsim,:,:])
Out[57]:
In [58]:
That,That_lin,That_nl = predict_kinetic(state.params, QnormTrainData[:Nsim,:,:])
Vhat,Vhat_lin,Vhat_nl = predict_potential(state.params, QnormTrainData[:Nsim,:,:])
Hhat = predict_Hamiltonian(state.params, QnormTrainData[:Nsim,:,:])
plt.figure(figsize = (16,12))
plt.subplot(3,1,1)
plt.plot(That_lin)
plt.plot(That_nl)
plt.title('kinetic energy')
plt.legend(['linear', 'nnet'])
plt.subplot(3,1,2)
plt.plot(Vhat_lin)
plt.plot(Vhat_nl)
plt.title('potential energy')
plt.legend(['linear', 'nnet'])
plt.subplot(3,1,3)
plt.plot(Vhat_lin + That_lin)
#plt.plot(Vhat + That)
plt.plot(Hhat[:,0,0], 'k--')
plt.title('total mechanical energy')
plt.show()
In [ ]:
In [60]:
plt.figure(figsize = (12,4))
plt.subplot(1,2,1)
plt.plot(That_lin - Vhat_lin)
plt.plot(That - Vhat, 'k--', alpha = 0.8)
ax = plt.gca()
ax2 = ax.twinx()
ax2.plot((That - Vhat).reshape((-1,)) - (That_lin - Vhat_lin).reshape((-1,)), 'm')
ax2.set_ylim([0, 0.5])
plt.title('Lagrangian (T - V)')
ax.legend(['linear', 'nnet'])
ax.set_ylabel('$L$')
ax2.set_ylabel('$L_{nn}$')
plt.subplot(1,2,2)
plt.plot(Vhat_lin + That_lin)
#plt.plot(Vhat + That)
plt.plot(Hhat[:,0,0], 'k--')
plt.legend(['linear', 'nnet'])
plt.ylabel('$H$')
plt.ylim([10,13])
plt.title('total mechanical energy (T + V)')
plt.tight_layout()
plt.show()
In [30]:
Mhat_pred,Khat_pred,Rhat_pred,Mhat_nl,Khat_nl,Chat_nl = predict_MKC(state.params, QnormTrainData[:Nsim,:,:])
In [61]:
plt.figure(figsize = (10,10))
labels = ['$\mathbf{M}$', '$\\delta \mathbf{M}(q)$','$\mathbf{K}$', '$\\delta \mathbf{K}(q)$','$\mathbf{C}$', '$\\delta \mathbf{C}(\dot{q})$']
plt.subplot(3,2,1)
plt.imshow(MHAT)
plt.colorbar()
plt.title(labels[0])
plt.axis('off')
plt.subplot(3,2,2)
plt.imshow(Mhat_nl[10])
plt.colorbar()
plt.title(labels[1])
plt.axis('off')
plt.subplot(3,2,3)
plt.imshow(KHAT)
plt.colorbar()
plt.title(labels[2])
plt.axis('off')
plt.subplot(3,2,4)
plt.imshow(Khat_nl[100])
plt.colorbar()
plt.title(labels[3])
plt.axis('off')
plt.subplot(3,2,5)
plt.imshow(CHAT)
plt.colorbar()
plt.title(labels[4])
plt.axis('off')
plt.subplot(3,2,6)
plt.imshow(Chat_nl[10])
plt.colorbar()
plt.title(labels[5])
plt.axis('off')
plt.tight_layout()
plt.show()
In [62]:
plt.figure(figsize = (10,10))
plt.subplot(2,2,1)
plt.plot(XnormTrainData[:Nsim,:,1])
plt.plot(np.diff(XnormTrainData[:Nsim,:,0], axis = 0)/dt, 'k--')
plt.subplot(2,2,2)
plt.plot(XnormTrainData[:Nsim,:,2])
plt.plot(np.diff(XnormTrainData[:Nsim,:,1], axis = 0)/dt, 'k--')
plt.subplot(2,2,3)
plt.plot(Xpred[:,:,1])
plt.plot(np.diff(Xpred[:,:,0], axis = 0)/dt, 'k--')
plt.subplot(2,2,4)
plt.plot(Xpred[:,:,2])
plt.plot(np.diff(Xpred[:,:,1], axis = 0)/dt, 'k--')
plt.show()
In [63]:
ddq_hat = predict_ddq(state.params, Qpred[:])
plt.figure()
plt.plot(ddq_hat)
plt.plot(Qpred[:,:,2], 'k--')
plt.show()
In [64]:
from scipy.signal import butter, lfilter, freqz
from scipy.interpolate import interp1d
Qext_pred = QnormTrainData[:,:,3]
print(Qext_pred.shape)
time_vec_int = t[:].reshape((-1,))
fbar_int = interp1d(time_vec_int[:Qext_pred.shape[0]], Qext_pred, axis = 0)
print(time_vec_int.shape, Qext_pred.shape)
In [65]:
qICs = np.concatenate([Qpred[0,:,0], Qpred[0,:,1]])
print(qICs)
In [71]:
import time
from scipy.interpolate import interp1d
import scipy.signal
dof = LATENT
def tf_dydx(q, dq, f, params):
return f_diff(params, q, dq, f)
def learned_q(x, t, params, fbar):
x = np.array(x)
t = np.array([t])
dof = LATENT
q = x[:dof]
dq = x[dof:]
f = fbar(t).reshape((-1,))
q = q.reshape((-1,LATENT))
dq = dq.reshape((-1,LATENT))
f = f.reshape((-1,LATENT))
dydx = tf_dydx(q, dq, f, params)
#qqdot = x.reshape((1, dof*2)) #np.concatenate([x, f], axis = 1)
#dydx = model.F_layers_q(qqdot).numpy()
return dydx.flatten()
Nint = Ntt*test_factor-10 #18000-10
import time
start_ = time.time()
qout = scipy.integrate.odeint(learned_q, qICs, time_vec_int[:Nint], (state.params,fbar_int))
end = time.time()
print(end - start_)
In [72]:
plt.figure(figsize = (10,10))
print(dof)
xpred = qout[:, :dof]
vpred = qout[:, dof:]
apred = f_diff(state.params, xpred, vpred, Qext_pred[:Nint,:])[:,dof:]
plt.subplot(3, 1, 1)
plt.plot(time_vec_int[:Nint], Qpred[:Nint,:,0], linewidth = 4, alpha = 0.6)
plt.plot(time_vec_int[:Nint], xpred, 'k')
plt.ylabel('Displacement')
plt.xlabel('Time, s')
plt.subplot(3, 1, 2)
plt.plot(time_vec_int[:Nint], Qpred[:Nint,:,1], linewidth = 4, alpha = 0.6)
plt.plot(time_vec_int[:Nint], vpred, 'k')
plt.ylabel('Velocity')
plt.xlabel('Time, s')
#plt.xlim([0, 0.1])
plt.subplot(3, 1, 3)
plt.plot(time_vec_int[:Nint], Qpred[:Nint,:,2], linewidth = 4, alpha = 0.6)
plt.plot(time_vec_int[:Nint], apred, 'k')
plt.ylabel('Acceleration')
plt.xlabel('Time, s')
#plt.xlim([0, 0.1])
plt.tight_layout()
#plt.savefig('latent_displacements_time_NL.tif')
#plt.savefig('/content/drive/MyDrive/UCSD_research/jax/chain_of_masses/time_integrated_latent_kprime2p0.png')
plt.show()
In [73]:
Qpred_int = np.concatenate([xpred[:,:,np.newaxis],
vpred[:,:,np.newaxis],
apred[:,:,np.newaxis],
Qext_pred[:Nint,:,np.newaxis]], axis = -1)
Xpred_int = decode_predict(state.params, Qpred_int)
print(Xpred_int.shape)
In [74]:
Xpred_Vq_q = np.dot(VMAT, Qpred_int[:,:,0].T).T
Xpred_Vq_dq = np.dot(VMAT, Qpred_int[:,:,1].T).T
Xpred_Vq = np.concatenate([Xpred_Vq_q[:,:,np.newaxis], Xpred_Vq_dq[:,:,np.newaxis]], axis = 2)
print(Xpred_Vq.shape)
In [75]:
np.random.seed(3)
plt.figure(figsize = (10,8))
ix = np.random.randint(0,64,3)
#start = 20
Nplot = Nint
plt.subplot(3, 1, 1)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,0], linewidth = 2, alpha = 0.6)
plt.plot(time_vec_int[:Nplot], Xpred_int[:Nplot,ix,0], 'k--')
#plt.vlines(time_vec_int[8000], ymin = -0.2, ymax = 0.2, color = 'r', linestyle = '--')
plt.ylabel('Displacement')
plt.xlabel('Time, s')
plt.title('NNet')
plt.subplot(3, 1, 2)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,0], linewidth = 2, alpha = 0.6)
plt.plot(time_vec_int[:Nplot], Qrr.T[start:Nplot+start,ix], 'k--')
#plt.vlines(time_vec_int[8000], ymin = -0.2, ymax = 0.2, color = 'r', linestyle = '--')
plt.ylabel('Displacement')
plt.xlabel('Time, s')
plt.title('ROM')
#errors_disp_nn = XnormTrainData[:Nplot,:,0] - Xpred_int[:Nplot,:,0]
#errors_disp_rom = XnormTrainData[:Nplot,:,0] - Qrr.T[start:Nplot+start,:]
#plt.subplot(3,1,3)
#for ii in range(64):
# plt.stem(time_vec_int[:Nplot], errors_disp_nn[:, ii], linefmt = 'blue')
# plt.stem(time_vec_int[:Nplot], errors_disp_rom[:, ii], linefmt = 'red')
plt.tight_layout()
plt.show()
In [78]:
np.random.seed(3)
plt.figure(figsize = (10,8))
ix = np.random.randint(0,64,3)
#start = 20
Nplot = Nint
plt.subplot(3, 1, 1)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,0], linewidth = 2, alpha = 0.6)
plt.plot(time_vec_int[:Nplot], Xpred_int[:Nplot,ix,0], 'k--')
plt.plot(time_vec_int[:Nplot], Xpred_Vq[:Nplot,ix,0], 'g--')
#plt.vlines(time_vec_int[8000], ymin = -0.2, ymax = 0.2, color = 'r', linestyle = '--')
plt.ylabel('Displacement')
plt.xlabel('Time, s')
plt.title('NNet')
plt.subplot(3, 1, 2)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,0], linewidth = 2, alpha = 0.6)
plt.plot(time_vec_int[:Nplot], Qrr.T[start:Nplot+start,ix], 'k--')
#plt.vlines(time_vec_int[8000], ymin = -0.2, ymax = 0.2, color = 'r', linestyle = '--')
plt.ylabel('Displacement')
plt.xlabel('Time, s')
plt.title('ROM')
#errors_disp_nn = XnormTrainData[:Nplot,:,0] - Xpred_int[:Nplot,:,0]
#errors_disp_rom = XnormTrainData[:Nplot,:,0] - Qrr.T[start:Nplot+start,:]
#plt.subplot(3,1,3)
#for ii in range(64):
# plt.stem(time_vec_int[:Nplot], errors_disp_nn[:, ii], linefmt = 'blue')
# plt.stem(time_vec_int[:Nplot], errors_disp_rom[:, ii], linefmt = 'red')
plt.tight_layout()
plt.show()
In [79]:
errors_disp_nn = np.sqrt(np.mean((XnormTrainData[:Nplot,:,0] - Xpred_int[:Nplot,:,0])**2))
errors_disp_rom = np.sqrt(np.mean((XnormTrainData[:Nplot,:,0] - Qrr.T[start:Nplot+start,:])**2))
errors_disp_nnVq = np.sqrt(np.mean((XnormTrainData[:Nplot,:,0] - Xpred_Vq[:Nplot,:,0])**2))
print(errors_disp_nn, errors_disp_rom, errors_disp_nnVq)
In [80]:
dQrr = np.diff(Qrr.T, axis = 0)/dt
dQrr = np.insert(dQrr, 0, XnormTrainData[0,:,1], axis = 0)
#dQrr = dQrr + XnormTrainData[0,:,1]
dQrr.shape
Out[80]:
In [81]:
np.random.seed(3)
plt.figure(figsize = (10,8))
ix = np.random.randint(0,64,3)
#start = 20
Nplot = Nint
plt.subplot(3, 1, 1)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,1], linewidth = 2, alpha = 0.6)
plt.plot(time_vec_int[:Nplot], Xpred_int[:Nplot,ix,1], 'k--')
#plt.vlines(time_vec_int[8000], ymin = -0.2, ymax = 0.2, color = 'r', linestyle = '--')
plt.ylabel('Velocity')
plt.xlabel('Time, s')
plt.title('NNet')
plt.subplot(3, 1, 2)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,1], linewidth = 2, alpha = 0.6)
plt.plot(time_vec_int[:Nplot], dQrr[start:Nplot+start,ix], 'k--')
#plt.vlines(time_vec_int[8000], ymin = -0.2, ymax = 0.2, color = 'r', linestyle = '--')
plt.ylabel('Velocity')
plt.xlabel('Time, s')
plt.title('ROM')
#errors_disp_nn = XnormTrainData[:Nplot,:,0] - Xpred_int[:Nplot,:,0]
#errors_disp_rom = XnormTrainData[:Nplot,:,0] - Qrr.T[start:Nplot+start,:]
#plt.subplot(3,1,3)
#for ii in range(64):
# plt.stem(time_vec_int[:Nplot], errors_disp_nn[:, ii], linefmt = 'blue')
# plt.stem(time_vec_int[:Nplot], errors_disp_rom[:, ii], linefmt = 'red')
plt.tight_layout()
plt.show()
In [82]:
np.random.seed(3)
plt.figure(figsize = (10,8))
ix = np.random.randint(0,64,3)
#start = 20
Nplot = Nint
plt.subplot(3, 1, 1)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,1], linewidth = 2, alpha = 0.6)
plt.plot(time_vec_int[:Nplot], Xpred_int[:Nplot,ix,1], 'k--')
plt.plot(time_vec_int[:Nplot], Xpred_Vq[:Nplot,ix,1], 'g--')
#plt.vlines(time_vec_int[8000], ymin = -0.2, ymax = 0.2, color = 'r', linestyle = '--')
plt.ylabel('Velocity')
plt.xlabel('Time, s')
plt.title('NNet')
plt.subplot(3, 1, 2)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,1], linewidth = 2, alpha = 0.6)
plt.plot(time_vec_int[:Nplot], dQrr[start:Nplot+start,ix], 'k--')
#plt.vlines(time_vec_int[8000], ymin = -0.2, ymax = 0.2, color = 'r', linestyle = '--')
plt.ylabel('Velocity')
plt.xlabel('Time, s')
plt.title('ROM')
#errors_disp_nn = XnormTrainData[:Nplot,:,0] - Xpred_int[:Nplot,:,0]
#errors_disp_rom = XnormTrainData[:Nplot,:,0] - Qrr.T[start:Nplot+start,:]
#plt.subplot(3,1,3)
#for ii in range(64):
# plt.stem(time_vec_int[:Nplot], errors_disp_nn[:, ii], linefmt = 'blue')
# plt.stem(time_vec_int[:Nplot], errors_disp_rom[:, ii], linefmt = 'red')
plt.tight_layout()
plt.show()
In [83]:
errors_vel_nn = np.sqrt(np.mean((XnormTrainData[:Nplot,:,1] - Xpred_int[:Nplot,:,1])**2))
errors_vel_rom = np.sqrt(np.mean((XnormTrainData[:Nplot,:,1] - dQrr[start:Nplot+start,:])**2))
errors_vel_nnVq = np.sqrt(np.mean((XnormTrainData[:Nplot,:,1] - Xpred_Vq[:Nplot,:,1])**2))
print(errors_vel_nn, errors_vel_rom, errors_vel_nnVq)
In [84]:
np.random.seed(3)
plt.figure(figsize = (10,8))
ix = np.random.randint(0,64,3)
ix = 2
start = 0
Nplot = Nint
plt.subplot(2, 2, 1)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,0], linewidth = 1, alpha = 0.7)
plt.plot(time_vec_int[:Nplot], Xpred_int[:Nplot,ix,0], 'r', alpha = 0.5)
plt.plot(time_vec_int[:Nplot], Xpred_Vq[:Nplot,ix,0], 'g--', alpha = 0.5)
plt.vlines(time_vec_int[Ntt], ymin = -0.06, ymax = 0.06, color = 'k', linestyle = '--')
plt.legend(['Truth', 'Pred mse (nl): {0:.4f}'.format(errors_disp_nn), 'Pred mse (lin): {0:.4f}'.format(errors_disp_nnVq)])
plt.ylabel('Displacement')
plt.xlabel('Time, s')
plt.title('NNet')
plt.subplot(2, 2, 2)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,0], linewidth = 1, alpha = 0.7)
plt.plot(time_vec_int[:Nplot], Qrr.T[start:Nplot+start,ix], 'r', alpha = 0.6)
plt.legend(['Truth', 'Pred mse: {0:.4f}'.format(errors_disp_rom)])
plt.vlines(time_vec_int[Ntt], ymin = -0.06, ymax = 0.06, color = 'k', linestyle = '--')
plt.ylabel('Displacement')
plt.xlabel('Time, s')
plt.title('ROM')
plt.subplot(2, 2, 3)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,1], linewidth = 1, alpha = 0.7)
plt.plot(time_vec_int[:Nplot], Xpred_int[:Nplot,ix,1], 'r', alpha = 0.5)
plt.plot(time_vec_int[:Nplot], Xpred_Vq[:Nplot,ix,1], 'g--', alpha = 0.5)
plt.legend(['Truth', 'Pred mse (nl): {0:.4f}'.format(errors_vel_nn), 'Pred mse (lin): {0:.4f}'.format(errors_vel_nnVq)])
plt.vlines(time_vec_int[Ntt], ymin = -0.6, ymax = 0.6, color = 'k', linestyle = '--')
plt.ylabel('Velocity')
plt.xlabel('Time, s')
plt.title('NNet')
plt.subplot(2, 2, 4)
plt.plot(time_vec_int[:Nplot], XnormTrainData[:Nplot,ix,1], linewidth = 1, alpha = 0.7)
plt.plot(time_vec_int[:Nplot], dQrr[start:Nplot+start,ix], 'r', alpha = 0.6)
plt.legend(['Truth', 'Pred mse: {0:.4f}'.format(errors_vel_rom)])
plt.vlines(time_vec_int[Ntt], ymin = -0.6, ymax = 0.6, color = 'k', linestyle = '--')
plt.ylabel('Velocity')
plt.xlabel('Time, s')
plt.title('ROM')
#errors_disp_nn = XnormTrainData[:Nplot,:,0] - Xpred_int[:Nplot,:,0]
#errors_disp_rom = XnormTrainData[:Nplot,:,0] - Qrr.T[start:Nplot+start,:]
#plt.subplot(3,1,3)
#for ii in range(64):
# plt.stem(time_vec_int[:Nplot], errors_disp_nn[:, ii], linefmt = 'blue')
# plt.stem(time_vec_int[:Nplot], errors_disp_rom[:, ii], linefmt = 'red')
plt.tight_layout()
plt.show()
In [85]:
np.random.seed(3)
plt.figure(figsize = (10,4))
ix = np.random.randint(0,64,1)
#start = 20
Nplot = Nint
plt.subplot(1, 2, 1)
plt.title('NNet')
plt.plot(XnormTrainData[:Nplot,ix,0], XnormTrainData[:Nplot,ix,1])
plt.plot(Xpred_int[:Nplot,ix,0], Xpred_int[:Nplot,ix,1], 'k')
plt.xlabel('$x$')
plt.ylabel('$\dot{x}$')
plt.legend(['Truth', 'Prediction'])
plt.subplot(1, 2, 2)
plt.title('ROM')
plt.plot(XnormTrainData[:Nplot,ix,0], XnormTrainData[:Nplot,ix,1])
plt.plot(Qrr.T[start:Nplot+start,ix], dQrr[start:Nplot+start,ix], 'k')
plt.xlabel('$x$')
plt.ylabel('$\dot{x}$')
plt.legend(['Truth', 'Prediction'])
plt.tight_layout()
plt.show()
In [86]:
np.random.seed(3)
plt.figure(figsize = (10,4))
ix = np.random.randint(0,64,1)
#start = 20
Nplot = Nint
plt.subplot(1, 2, 1)
plt.title('NNet')
plt.plot(XnormTrainData[:Nplot,ix,0], XnormTrainData[:Nplot,ix,1])
plt.plot(Xpred_int[:Nplot,ix,0], Xpred_int[:Nplot,ix,1], 'k')
plt.plot(Xpred_Vq[:Nplot,ix,0], Xpred_Vq[:Nplot,ix,1], 'g--')
plt.xlabel('$x$')
plt.ylabel('$\dot{x}$')
#plt.legend(['Truth', 'Prediction', 'Linear Project NNet'])
plt.subplot(1, 2, 2)
plt.title('ROM')
plt.plot(XnormTrainData[:Nplot,ix,0], XnormTrainData[:Nplot,ix,1])
plt.plot(Qrr.T[start:Nplot+start,ix], dQrr[start:Nplot+start,ix], 'k')
plt.xlabel('$x$')
plt.ylabel('$\dot{x}$')
plt.legend(['Truth', 'Prediction'])
plt.tight_layout()
plt.show()
In [87]:
fig = plt.figure(figsize = (16,14))
np.random.seed(190)
ix = np.random.randint(0,64,3)
Nplot = Ntt*2-10
ax = fig.add_subplot(1, 2, 1, projection='3d')
ax.plot3D(XnormTrainData[:Nplot,ix[0],0], XnormTrainData[:Nplot,ix[1],0], XnormTrainData[:Nplot,ix[0],1], linewidth = 2)
ax.plot3D(Xpred_int[:Nplot,ix[0],0], Xpred_int[:Nplot,ix[1],0], Xpred_int[:Nplot,ix[0],1], 'k--', linewidth = 1)
#ax.plot3D(delta_x_hat[tmin:Nint], delta_xdot_hat[tmin:Nint], delta_Q_hat[tmin:Nint], 'k--')
plt.xlabel('$x_{{{}}}$'.format(ix[0]), fontsize = 28)
plt.ylabel('$x_{{{}}}$'.format(ix[1]), fontsize = 28)
ax.set_zlabel('$\dot{{x}}_{{{}}}$'.format(ix[0]), fontsize = 28)
ax = fig.add_subplot(1, 2, 2, projection='3d')
ax.plot3D(XnormTrainData[:Nplot,ix[0],0], XnormTrainData[:Nplot,ix[1],0], XnormTrainData[:Nplot,ix[0],1], linewidth = 2)
ax.plot3D(Qrr.T[:Nplot,ix[0]], Qrr.T[:Nplot,ix[1]], dQrr[:Nplot,ix[0]], 'k--', linewidth = 1)
plt.xlabel('$x_{{{}}}$'.format(ix[0]), fontsize = 28)
plt.ylabel('$x_{{{}}}$'.format(ix[1]), fontsize = 28)
ax.set_zlabel('$\dot{{x}}_{{{}}}$'.format(ix[0]), fontsize = 28)
plt.tight_layout()
plt.show()
In [88]:
fig = plt.figure(figsize = (16,14))
np.random.seed(190)
ix = np.random.randint(0,64,3)
Nplot = Ntt*2-10
ax = fig.add_subplot(1, 2, 1, projection='3d')
ax.plot3D(XnormTrainData[:Nplot,ix[0],0], XnormTrainData[:Nplot,ix[1],0], XnormTrainData[:Nplot,ix[0],1], linewidth = 2)
ax.plot3D(Xpred_int[:Nplot,ix[0],0], Xpred_int[:Nplot,ix[1],0], Xpred_int[:Nplot,ix[0],1], 'k--', linewidth = 1)
ax.plot3D(Xpred_Vq[:Nplot,ix[0],0], Xpred_Vq[:Nplot,ix[1],0], Xpred_Vq[:Nplot,ix[0],1], 'g--', linewidth = 1)
#ax.plot3D(delta_x_hat[tmin:Nint], delta_xdot_hat[tmin:Nint], delta_Q_hat[tmin:Nint], 'k--')
plt.xlabel('$x_{{{}}}$'.format(ix[0]), fontsize = 28)
plt.ylabel('$x_{{{}}}$'.format(ix[1]), fontsize = 28)
ax.set_zlabel('$\dot{{x}}_{{{}}}$'.format(ix[0]), fontsize = 28)
ax = fig.add_subplot(1, 2, 2, projection='3d')
ax.plot3D(XnormTrainData[:Nplot,ix[0],0], XnormTrainData[:Nplot,ix[1],0], XnormTrainData[:Nplot,ix[0],1], linewidth = 2)
ax.plot3D(Qrr.T[:Nplot,ix[0]], Qrr.T[:Nplot,ix[1]], dQrr[:Nplot,ix[0]], 'k--', linewidth = 1)
plt.xlabel('$x_{{{}}}$'.format(ix[0]), fontsize = 28)
plt.ylabel('$x_{{{}}}$'.format(ix[1]), fontsize = 28)
ax.set_zlabel('$\dot{{x}}_{{{}}}$'.format(ix[0]), fontsize = 28)
plt.tight_layout()
plt.show()
In [ ]: